
#include "lerp_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <algorithm>


namespace optiling {
const uint32_t BLOCK_SIZE = 32;
const uint32_t BUFFER_NUM = 2;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    LerpTilingData tiling;
    
    // uint32_t sizeofdatatype;
    uint64_t ubSize;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    // auto socVersion = ascendcPlatform.GetSocVersion();
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); //获取硬件平台存储空间 UB 的内存大小
    // auto aivNum = ascendcPlatform.GetCoreNum(); //获取当前硬件平台的核数 此平台为1

    //获取输入shape信息
    uint32_t inputNum = context->GetInputShape(1)->GetStorageShape().GetShapeSize();
    uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(0)->GetDataType());
    uint32_t inputLength = inputBytes * inputNum;

    //可使用的ub空间 输入3输出1
    uint32_t ubDataNumber = (inputBytes == 2) ? 10 : 4;

    // The number of 32B data blocks that can be used for each data. DOUBLE BUFFER is already counted here
    uint32_t tileBlockNum = (ubSize / BLOCK_SIZE / BUFFER_NUM) / ubDataNumber;
    uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes;

    // Input data for 32B alignment
    uint32_t inputLengthAlgin32 = (((inputLength + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE);
    // There is at least 32B of data on each core, satisfying several settings for several cores. The maximum number of audits is the actual number of audits
    uint32_t everyCoreInputBlockNum = inputLengthAlgin32 / BLOCK_SIZE;// aivNum;
    
    //  chunks are calculated and sliced several times using the number of data on each core
    uint32_t CoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
    uint32_t TileNum = everyCoreInputBlockNum / tileBlockNum;
    uint32_t finalTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? TileNum : TileNum + 1;
    // Tail block calculation for  chunks of data
    uint32_t TailDataNum = CoreDataNum - (tileDataNum * TileNum);
    TailDataNum = TailDataNum == 0 ? tileDataNum : TailDataNum;

    tiling.set_CoreDataNum(CoreDataNum);
    tiling.set_tileDataNum(tileDataNum);
    tiling.set_TailDataNum(TailDataNum);
    tiling.set_finalTileNum(finalTileNum);

    uint32_t weightSize = context->GetInputShape(2)->GetStorageShape().GetShapeSize();
    uint32_t startSize = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
    uint32_t endSize = context->GetInputShape(1)->GetStorageShape().GetShapeSize();
    uint32_t ySize = context->GetOutputShape(0)->GetStorageShape().GetShapeSize();
    
    if(ySize != startSize || ySize != endSize || ySize != weightSize)
    {
        context->SetTilingKey(2);

        int32_t y_ndarray[20], start_ndarray[20], end_ndarray[20], weight_ndarray[20];
        int32_t y_dimensional, start_dimensional, end_dimensional, weight_dimensional;
        auto shape_y = context->GetOutputShape(0)->GetOriginShape();
        auto shape_start = context->GetInputTensor(0)->GetOriginShape();
        auto shape_end = context->GetInputTensor(1)->GetOriginShape();
        auto shape_weight = context->GetInputTensor(2)->GetOriginShape();

        y_dimensional =  shape_y.GetDimNum();
        start_dimensional =  shape_start.GetDimNum();
        end_dimensional =  shape_end.GetDimNum();
        weight_dimensional =  shape_weight.GetDimNum();

        for(int i = 0; i < y_dimensional; i++)
        {
            y_ndarray[y_dimensional-i-1] = shape_y.GetDim(i);
            if(i<start_dimensional) start_ndarray[start_dimensional-i-1] = shape_start.GetDim(i);
            else                    start_ndarray[i] = 1;
            if(i<end_dimensional) end_ndarray[end_dimensional-i-1] = shape_end.GetDim(i);
            else                  end_ndarray[i] = 1;
            if(i<weight_dimensional) weight_ndarray[weight_dimensional-i-1] = shape_weight.GetDim(i);
            else                     weight_ndarray[i] = 1;
        }
        
        tiling.set_y_dimensional(y_dimensional);
        tiling.set_y_ndarray(y_ndarray);
        tiling.set_start_ndarray(start_ndarray);
        tiling.set_end_ndarray(end_ndarray);
        tiling.set_weight_ndarray(weight_ndarray);

        int32_t y_sumndarray[20], start_sumndarray[20], end_sumndarray[20], weight_sumndarray[20];
        y_sumndarray[0] = 1;
        start_sumndarray[0] = 1;
        end_sumndarray[0] = 1;
        weight_sumndarray[0] = 1;
        for(int i = 1; i <= y_dimensional; i++)
        {
            y_sumndarray[i] = y_sumndarray[i-1]*y_ndarray[i-1];
            start_sumndarray[i] = start_sumndarray[i-1]*start_ndarray[i-1];
            end_sumndarray[i] = end_sumndarray[i-1]*end_ndarray[i-1];
            weight_sumndarray[i] = weight_sumndarray[i-1]*weight_ndarray[i-1];
        }
        tiling.set_y_sumndarray(y_sumndarray);
        tiling.set_start_sumndarray(start_sumndarray);
        tiling.set_end_sumndarray(end_sumndarray);
        tiling.set_weight_sumndarray(weight_sumndarray);
    }
    else
    {
        context->SetTilingKey(1);
    }

    context->SetBlockDim(1);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Lerp : public OpDef {
public:
    explicit Lerp(const char* name) : OpDef(name)
    {
        this->Input("start")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("end")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("weight")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(Lerp);
}
